#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist, TransformStamped
from tf2_ros import TransformListener, Buffer
import math

class ArucoFollower(Node):
    def __init__(self):
        super().__init__('aruco_follower')
        self.cmd_pub = self.create_publisher(Twist, '/cmd_vel', 10)
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        self.timer = self.create_timer(0.1, self.control_loop)
        
        # 控制参数
        self.target_distance = 0.5  # 目标跟随距离（米）
        self.kp_linear = 0.5       # 线性速度比例系数
        self.kp_angular = 1.0      # 角速度比例系数

    def control_loop(self):
        try:
            # 获取相机到标记的TF变换
            trans = self.tf_buffer.lookup_transform(
                'camera_link', 'aruco_marker', rclpy.time.Time())
            
            # 计算距离和角度误差
            distance = math.sqrt(trans.transform.translation.x**2 + 
                                trans.transform.translation.z**2)
            angle = math.atan2(trans.transform.translation.x,
                              trans.transform.translation.z)
            
            # PID控制（简化版）
            cmd = Twist()
            if distance > self.target_distance * 1.2:
                cmd.linear.x = min(0.5, self.kp_linear * (distance - self.target_distance))
            elif distance < self.target_distance * 0.8:
                cmd.linear.x = max(-0.5, -self.kp_linear * (self.target_distance - distance))
            
            cmd.angular.z = -self.kp_angular * angle
            
            self.cmd_pub.publish(cmd)
            
        except Exception as e:
            self.get_logger().warn(f'TF error: {str(e)}')
            # 丢失标记时停止
            self.cmd_pub.publish(Twist())

def main():
    rclpy.init()
    node = ArucoFollower()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()